from typing import Annotated
from typing_extensions import TypedDict
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from utils.logger_config import LoggerSingleton
logger = LoggerSingleton.get_logger() 

# 定义整个对话的状态
class State(TypedDict):
    # messages的类型是"list"。
    # "add_messages"函数定义了如何更新此状态键
    # （在这里，它将messages追加到列表中，而不是覆盖它们）。
    messages: Annotated[list, add_messages]

# 定义初始图
logger.info('graph初始化: 增加模型节点、增加边、编译')
graph_builder = StateGraph(State)

# 定义模型节点函数
# 添加模型节点
## 导入模型

from model_init.get_model_ollama import ollama_model
logger.info(f'模型name: {ollama_model.model}')
## 创建节点函数
def chatbot(state: State):
    return {"messages": [ollama_model.invoke(state["messages"])]}

## 将节点加入到图中
### 第一个参数是唯一的节点名。
### 第二个参数是当节点被使用时调用的函数或对象。
graph_builder.add_node("chatbot", chatbot)
# 定义连接边
graph_builder.add_edge(START, "chatbot")
graph_builder.add_edge("chatbot", END)
# 编译为可执行状态
graph = graph_builder.compile()

if __name__ == "__main__":
    # 绘图——可视化
    flag_visual = 1
    if flag_visual != 0:
        logger.info('graph可视化')
        from utils.visual_drawing import save_graph_image
        save_graph_image(graph)
    else:
        logger.info('graph未可视化')

    def stream_graph_updates(user_input: str):
        for event in graph.stream({"messages": [{"role": "user", "content": user_input}]}):
            for value in event.values():
                print("Assistant:", value["messages"][-1].content)

    while True:
        try:
            user_input = input("User: ")
            if user_input.lower() in ["quit", "exit", "q"]:
                print("Goodbye!")
                break

            stream_graph_updates(user_input)
        except:
            # fallback if input() is not available
            user_input = "What do you know about LangGraph?"
            print("User: " + user_input)
            stream_graph_updates(user_input)
            break